###################################
#
#           相机标定程序
#
###################################
import cv2
import numpy as np
import yaml
import os

# 棋盘格尺寸和边长（单位：cm）
CHESSBOARD_SIZE = (8,11)  # 棋盘格内部角点数
SQUARE_SIZE = 2.0          # 每个方格的边长（cm）

# 存储有效帧的采样率
SAMPLING_RATE = 0.1

# 创建对象点（世界坐标系中的角点坐标）
objp = np.zeros((CHESSBOARD_SIZE[0] * CHESSBOARD_SIZE[1], 3), np.float32)
objp[:, :2] = np.mgrid[0:CHESSBOARD_SIZE[0], 0:CHESSBOARD_SIZE[1]].T.reshape(-1, 2) * SQUARE_SIZE

# 存储角点数据
objpoints = []  # 世界坐标系中的角点
imgpoints = []  # 图像坐标系中的角点

# 打开相机
cap = cv2.VideoCapture(0)

if not cap.isOpened():
    print("无法打开相机")
    exit()

print("按 Esc 键退出并开始标定")

frame_count = 0
while True:
    ret, frame = cap.read()
    if not ret:
        print("无法读取相机帧")
        break

    # 转换为灰度图像
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 查找棋盘格角点
    ret_corners, corners = cv2.findChessboardCorners(gray, CHESSBOARD_SIZE, None)

    # 如果找到角点
    if ret_corners:
        # 绘制角点
        cv2.drawChessboardCorners(frame, CHESSBOARD_SIZE, corners, ret_corners)

        # 按照采样率存储有效帧
        if frame_count % int(1 / SAMPLING_RATE) == 0:
            objpoints.append(objp)
            imgpoints.append(corners)
            print(f"已存储 {len(objpoints)} 帧")

    # 显示图像
    cv2.imshow("Camera Calibration", frame)

    # 按下 Esc 键退出
    key = cv2.waitKey(1)
    if key == 27:  # Esc 键的 ASCII 码为 27
        break

    frame_count += 1

# 释放相机资源
cap.release()
cv2.destroyAllWindows()

# 如果有足够的有效帧，进行标定
if len(objpoints) > 0:
    print("开始标定相机...")
    ret, camera_matrix, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(
        objpoints, imgpoints, gray.shape[::-1], None, None
    )

    if ret:
        print("标定成功！")
        print("相机内参矩阵：\n", camera_matrix)
        print("畸变系数：\n", dist_coeffs)

        # 将标定参数保存到 YAML 文件
        calibration_data = {
            "camera_matrix": camera_matrix.tolist(),
            "distortion_coefficients": dist_coeffs.tolist(),
        }

        with open("camera_param.yaml", "w") as f:
            yaml.dump(calibration_data, f)
        print("标定参数已保存到 camera_param.yaml")
    else:
        print("标定失败")
else:
    print("没有足够的有效帧进行标定")